from classfier.TotalModel import *

class ModelUtil:
    def __init__(self, dp,X,split_random_state,save_path,model_random_state,data_type,model_type='logic',feature_extract_method='lasso'):
        self.dp = dp
        self.X = X
        self.split_random_state = split_random_state
        self.model_random_state = model_random_state
        self.save_path = save_path
        self.model_type = model_type
        self.data_type = data_type
        self.feature_extract_method = feature_extract_method

    def get_model(self):
        model_map = {
            'logic': LogisticUtil,
            'svm': SVMUtil,
            'rf': RFUtil,
            'gbm': GBMUtil,
            'nn': NNUtil,
            'xgb': XGBoostUtil,
            'knn': KNNUtil,
            'ada': AdaBoostUtil,
            'cat': CatBoostUtil
        }

        if self.model_type not in model_map:
            raise ValueError(f"不支持的模型类型：{self.model_type}")

        return model_map[self.model_type](
            dp=self.dp,
            X=self.X,
            split_random_state=self.split_random_state,
            model_random_state=self.model_random_state,
            save_path=self.save_path,
            data_type=self.data_type
        )

